#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
ComfyUI LoRA训练插件
"""

import os
import sys
import importlib

# 添加必要的路径
current_dir = os.path.dirname(os.path.abspath(__file__))
sd_scripts_path = os.path.join(current_dir, "sd_scripts")

for path in [current_dir, sd_scripts_path]:
    if os.path.exists(path) and path not in sys.path:
        sys.path.insert(0, path)

# 从constants.py导入训练器专用参数类型
from .nodes.sd_train.lora_trainer_utils.constants import (
    FLUX_TRAINING_PARAMS,
    SD3_TRAINING_PARAMS,
    SDXL_TRAINING_PARAMS,
    SD_TRAINING_PARAMS,
    BASIC_TRAINING_PARAMS
)

# 检查并修复huggingface_hub兼容性
try:
    # 确保nodes目录在sys.path中
    nodes_path = os.path.join(current_dir, "nodes")
    if nodes_path not in sys.path:
        sys.path.insert(0, nodes_path)
    
    # 现在应该可以导入了
    from fix import fix_huggingface_hub
    if fix_huggingface_hub.run_fix():
        print("[INFO] huggingface_hub兼容性检测/修复已完成")
    else:
        print("[WARNING] huggingface_hub兼容性修复失败，请手动处理！")
except Exception as e:
    print(f"[WARNING] huggingface_hub兼容性检测/修复时出错: {e}")

# 节点映射字典
NODE_CLASS_MAPPINGS = {}
NODE_DISPLAY_NAME_MAPPINGS = {}

def load_all_nodes():
    """加载nodes目录下所有子目录中的节点"""
    nodes_dir = os.path.join(current_dir, "nodes")
    if not os.path.exists(nodes_dir):
        print(f"[WARNING] nodes目录不存在: {nodes_dir}")
        return
    
    print("[INFO] 开始加载节点...")
    
    # 遍历nodes目录下的所有子目录
    for subdir in os.listdir(nodes_dir):
        subdir_path = os.path.join(nodes_dir, subdir)
        
        # 跳过非目录和特殊目录
        if not os.path.isdir(subdir_path) or subdir.startswith('__') or subdir == '__pycache__':
            continue
            
        print(f"[INFO] 加载 {subdir} 节点...")
        
        # 遍历子目录中的所有.py文件
        for file_name in os.listdir(subdir_path):
            if file_name.endswith('.py') and not file_name.startswith('__'):
                module_name = os.path.splitext(file_name)[0]
                try:
                    # 导入模块
                    module = importlib.import_module(f".nodes.{subdir}.{module_name}", package=__package__)
                    
                    # 收集映射
                    if hasattr(module, 'NODE_CLASS_MAPPINGS'):
                        NODE_CLASS_MAPPINGS.update(module.NODE_CLASS_MAPPINGS)
                    if hasattr(module, 'NODE_DISPLAY_NAME_MAPPINGS'):
                        NODE_DISPLAY_NAME_MAPPINGS.update(module.NODE_DISPLAY_NAME_MAPPINGS)
                        
                except ImportError as e:
                    print(f"[WARNING] 无法导入模块 {subdir}.{module_name}: {e}")
                except Exception as e:
                    print(f"[ERROR] 导入模块 {subdir}.{module_name} 时出错: {e}")

# 加载所有节点
load_all_nodes()

# 显式导入并初始化系统监控API
try:
    from .nodes.system_monitor.system_monitor_api import initialize_system_monitor
    initialize_system_monitor()
except ImportError as e:
    print(f"[ERROR] 无法导入系统监控API: {e}")
except Exception as e:
    print(f"[ERROR] 初始化系统监控API时出错: {e}")

print(f"[INFO] 节点加载完成，共加载 {len(NODE_CLASS_MAPPINGS)} 个节点")

# Web扩展支持
WEB_DIRECTORY = "./web/js"

__all__ = [
    "NODE_CLASS_MAPPINGS", 
    "NODE_DISPLAY_NAME_MAPPINGS", 
    "WEB_DIRECTORY",
    "FLUX_TRAINING_PARAMS",
    "SD3_TRAINING_PARAMS", 
    "SDXL_TRAINING_PARAMS",
    "SD_TRAINING_PARAMS",
    "BASIC_TRAINING_PARAMS"
]

__version__ = "0.1.0"
